/* -----------------------------------------------------------------------
   See COPYRIGHT.TXT and LICENSE.TXT for copyright and license information
   ----------------------------------------------------------------------- */
#include "plm_config.h"
#include <float.h>
#include <math.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
//#include "bspline.h"
//#include "bspline_optimize.h"
#include "pvalue_optimize_lbfgsb.h"
//#include "pvalue_mse.h"
#include "bspline_opts.h"
#include "logfile.h"
#include "plm_fortran.h"
#include "volume.h"
#include "print_and_exit.h"
#include "proj_image.h"
#include "proj_image_set.h"
#include "CB_set_op.h"
#include "drr_opts.h"
#include "zhang_cuda.h"
#include "fdk_util.h"
#include "mha_io.h"

#if defined __cplusplus
extern "C" {
#endif
void
setulb_ (integer*       n,
	 integer*       m,
	 doublereal*    x,
	 doublereal*    l,
	 doublereal*    u,
	 integer*       nbd,
	 doublereal*    f,
	 doublereal*    g,
	 doublereal*    factr,
	 doublereal*    pgtol,
	 doublereal*    wa,
	 integer*       iwa,
	 char*          task,
	 integer*       iprint,
	 char*          csave,
	 logical*       lsave,
	 integer*       isave,
	 doublereal*    dsave,
	 ftnlen         task_len,
	 ftnlen         csave_len
	 );
#if defined __cplusplus
}
#endif



void
pvalue_score_mse (    Bspline_parms *parms, 
				  Bspline_state *bst,
				  Volume *Vx, 
				  Volume *Vg,
				  CB_Set *PSd,
				  CB_Set *PSx,
				  void * fdk_dev_state,
				  void * drr_dev_state,
				  Drr_options *drr_options
				  ) 
{
	 Bspline_score* ssd = &bst->ssd;

	 CUDA_Pj(PSx,Vx, drr_dev_state, drr_options);
	 CB_Set_diff(PSx,PSd);
	 ssd->score=CB_Set_nrm(PSx);
	 CUDA_Bkpj(Vg,PSx,fdk_dev_state);
	/* for (int i=0; i<ssd->num_vox; i++)
		 ssd->grad[i]=((float *)Vg->img )[i];*/
}


class Nocedal_optimizer
{
public:
    char task[60], csave[60];
    logical lsave[4];
    integer n, m, iprint, *nbd, *iwa, isave[44];
    doublereal f, factr, pgtol, *x, *l, *u, *g, *wa, dsave[29];
public:
    Nocedal_optimizer (Volume *vol, Bspline_parms * parms);
    ~Nocedal_optimizer () {
	free (nbd);
	free (iwa);
	free (x);
	free (l);
	free (u);
	free (g);
	free (wa);
    }
    void setulb () {
	setulb_ (&n,&m,x,l,u,nbd,&f,g,&factr,&pgtol,wa,iwa,task,&iprint,
	    csave,lsave,isave,dsave,60,60);
    }
};

Nocedal_optimizer::Nocedal_optimizer (Volume *vol, Bspline_parms * parms)
{
   // Bspline_xform *bxf = bod->bxf;
   //Bspline_parms *parms = bod->parms;

    int NMAX = vol->npix;
    int MMAX = 20;

    /* Try to allocate memory for hessian approximation.  
       First guess based on heuristic. */
    if (vol->npix >= 20) {
	MMAX = 20 + (int) floor (sqrt ((float) (vol->npix - 20)));
	if (MMAX > 1000) {
	    MMAX = 1000;
	}
    }
    do {
	nbd = (integer*) malloc (sizeof(integer)*NMAX);
	iwa = (integer*) malloc (sizeof(integer)*3*NMAX);
	x = (doublereal*) malloc (sizeof(doublereal)*NMAX);
	l = (doublereal*) malloc (sizeof(doublereal)*NMAX);
	u = (doublereal*) malloc (sizeof(doublereal)*NMAX);
	g = (doublereal*) malloc (sizeof(doublereal)*NMAX);
	wa = (doublereal*) malloc (sizeof(doublereal)
	    *(2*MMAX*NMAX+4*NMAX+12*MMAX*MMAX+12*MMAX));

	if ((nbd == NULL) ||
	    (iwa == NULL) ||
	    (  x == NULL) ||
	    (  l == NULL) ||
	    (  u == NULL) ||
	    (  g == NULL) ||
	    ( wa == NULL))
	{
	    /* We didn't get enough memory.  Free what we got. */
	    free (nbd);
	    free (iwa);
	    free (x);
	    free (l);
	    free (u);
	    free (g);
	    free (wa);

	    /* Give a little feedback to the user */
	    logfile_printf (
		"Tried NMAX, MMAX = %d %d, but ran out of memory!\n",
		NMAX, MMAX);

	    /* Try again with reduced request */
	    if (MMAX > 20) {
		MMAX = MMAX / 2;
	    } else if (MMAX > 10) {
		MMAX = 10;
	    } else if (MMAX > 2) {
		MMAX --;
	    } else {
		print_and_exit ("System ran out of memory when "
		    "initializing Nocedal optimizer.\n");
	    }
	}
	else {
	    /* Everything went great.  We got the memory. */
	    break;
	}
    } while (1);
    m = MMAX;
    n = NMAX;

    /* Give a little feedback to the user */
    logfile_printf ("Setting NMAX, MMAX = %d %d\n", NMAX, MMAX);

    /* If iprint is 1, the file iterate.dat will be created */
    iprint = 0;

    //factr = 1.0e+7;
    //pgtol = 1.0e-5;
    factr = parms->lbfgsb_factr;
    pgtol = parms->lbfgsb_pgtol;

    /* Bounds for deformation problem */
    for (int i = 0; i < n; i++) {
	nbd[i] = 0;
	l[i]=-1.0e1;
	u[i]=+1.0e1;
    }

    /* Initial guess */
    for (int i = 0; i < n; i++) {
	x[i] = 0;
    }

    /* Remember: Fortran expects strings to be padded with blanks */
    memset (task, ' ', sizeof(task));
    memcpy (task, "START", 5);
    logfile_printf (">>> %c%c%c%c%c%c%c%c%c%c\n", 
	task[0], task[1], task[2], task[3], task[4], 
	task[5], task[6], task[7], task[8], task[9]);
}

void
pvalue_optimize_lbfgsb (
    Volume * vol, CB_Set * PSd, Bspline_parms * parms, Bspline_state *bst, void * fdk_dev_state, void * drr_dev_state, Fdk_options *cb_options, Drr_options * drr_options
	)
{
    //Bspline_xform *bxf = bod->bxf;
    //Bspline_state *bst = bod->bst;
    //Bspline_parms *parms = bod->parms;
   // Volume *fixed = bod->fixed;
    //Volume *moving = bod->moving;
    //Volume *moving_grad = bod->moving_grad;
    Bspline_score* ssd = &bst->ssd;
    FILE *fp = 0;
    double best_score = DBL_MAX;
    float *best_coeff = (float*) malloc (sizeof(float) * vol->npix);

    Nocedal_optimizer optimizer (vol, parms);

	Volume *Vg=my_create_volume(cb_options);

	CB_Set PSx;
	PSx.nproj=PSd->nproj;
	for (int i=0; i<PSx.nproj; i++){
		PSx.cbi[i]=proj_image_create();
		proj_image_create_img(PSx.cbi[i],PSd->cbi[i]->dim);
		memset(PSx.cbi[i]->img,0,sizeof(float)*PSd->cbi[i]->dim[0]*PSd->cbi[i]->dim[1]);
		PSx.cbi[i]->pmat=proj_matrix_clone(PSd->cbi[i]->pmat);
	}

    /* Initialize # iterations, # function evaluations */
   bst->it = 0;
   bst->feval = 0;

    if (parms->debug) {
	fp = fopen ("scores.txt", "w");
    }

    while (1) {
	/* Get next search location */
	optimizer.setulb ();

	if (optimizer.task[0] == 'F' && optimizer.task[1] == 'G') {
	    /* Got a new probe location within a line search */

	    /* Copy from fortran to C (double -> float) */
	    for (int i = 0; i < vol->npix; i++) {
		((float*)vol->img)[i] = (float) optimizer.x[i];
	    }

	    /* Compute cost and gradient */

		
	    pvalue_score_mse (parms, bst, vol,Vg, PSd, &PSx, fdk_dev_state, drr_dev_state, drr_options);

		write_mha("C:\\AAAFiles\\CBCTData\\Vg.mh5",Vg);
		//write_mha("C:\\AAAFiles\\CBCTData\\vol.m5",vol);



		//pvalue_score_mse (    Bspline_parms *parms, 
		//	Bspline_state *bst,
		//	Volume *Vx, 
		//	Volume *Vg,
		//	CB_Set *PSd,
		//	CB_Set *PSx,
		//	void * fdk_dev_state,
		//	void * drr_dev_state,
		//	Drr_options *drr_options
		//	) 

	    /* Save coeff if best score */
	    if (ssd->score < best_score) {
		best_score = ssd->score;
		for (int i = 0; i < vol->npix; i++) {
		    best_coeff[i] = ((float*)vol->img)[i];
		}
	    }

	    /* Give a little feedback to the user */
	    //bspline_display_coeff_stats (bxf);

	    /* Save some debugging information */
	    //bspline_save_debug_state (parms, bst, bxf);
	    if (parms->debug) {
		fprintf (fp, "%f\n", ssd->score);
	    }

	    /* Copy from C to fortran (float -> double) */
	    optimizer.f = ssd->score;
	    for (int i = 0; i < vol->npix; i++) {
		optimizer.g[i] =((float*)Vg->img)[i];
	    }

	    /* Check # feval */
	    if (bst->feval >= parms->max_feval) break;
	    bst->feval ++;

	} else if (memcmp (optimizer.task, "NEW_X", strlen ("NEW_X")) == 0) {
	    /* Optimizer has completed a line search. */

	    /* Check iterations */
	    if (bst->it >= parms->max_its) break;
	    bst->it ++;

	} else {
	    break;
	}
    }

    if (parms->debug) {
	fclose (fp);
    }

    /* Copy out the best results */
    for (int i = 0; i < vol->npix; i++) {
	((float*)vol->img)[i] = best_coeff[i];
    }
    free (best_coeff);
}



